import numpy as np
from numpy.typing import NDArray
from typing import Dict, Tuple

from dataset_sdk import PointXYZ, Obstacle, ObstacleCategory

__all__ = ['ObjectTrack']

def diag_stack(*np_mats: NDArray):
    dim0 = dim1 = 0
    str_pos = [[0, 0]]
    for mat in np_mats:
        assert mat.ndim == 2
        dim0 += mat.shape[0]
        dim1 += mat.shape[1]
        str_pos.append([dim0, dim1])
    res = np.zeros((dim0, dim1))
    for i, mat in enumerate(np_mats):
        res[str_pos[i][0]:str_pos[i][0] + mat.shape[0], str_pos[i][1]:str_pos[i][1] + mat.shape[1]] = mat
    return res

INVALID_MOTION = [-1000.] * 3
STAMP_PER_SEC = 1e6

class ObjectTrack():
    def __init__(self, timestamp_ms:int, obj: Obstacle) -> None:
        self.traces: Dict[int, Obstacle] = {timestamp_ms: obj}
        self.velocity: Dict[int, PointXYZ] = {}
        self.acceleration: Dict[int, PointXYZ] = {}

    def get_id(self):
        assert len(self.traces) > 0
        return list(self.traces.values())[0].id

    def get_category(self) -> ObstacleCategory:
        assert len(self.traces) > 0
        return list(self.traces.values())[0].category

    def infer_motion(self, ca_interval_s: float) -> None:
        def debug_svd(i: int) -> None:
            mat_t = mat_A_list[i]
            mat_p = mat_b_list[i]
            motionx = motion_mat[i*2:i*2+2, i*2]
            motiony = motion_mat[i*2:i*2+2, i*2+1]
            import matplotlib.pyplot as plt
            xdots = []
            ydots = []
            for j in range(mat_t.shape[0]):
                xdots.append(mat_p[j, 0])
                ydots.append(mat_p[j, 1])
            _, axs = plt.subplots(ncols=2)
            axs[0].scatter(mat_t[:, 0], np.array(xdots), c='r')
            axs[1].scatter(mat_t[:, 0], np.array(ydots), c='r')
            plot_t = np.arange(min(mat_t[:, 0].min(), 0), max(mat_t[:, 0].max(), 0), 0.2)
            plot_t = np.vstack([plot_t, plot_t ** 2 / 2]).T
            axs[0].plot(plot_t[:, 0], plot_t.dot(motionx))
            axs[1].plot(plot_t[:, 0], plot_t.dot(motiony))
            stamp = list(reference_map.keys())[i]
            fig = plt.figure(f'trace{i} ref boxes')
            self.traces[stamp].matplot_on_bev(axes=fig.gca())
            for ref_t in reference_map[stamp]:
                self.traces[ref_t].matplot_on_bev(rgb='orange', axes=fig.gca())

        ca_interval_us_int = int(ca_interval_s * 5e5)
        frame_stamps = list(self.traces.keys())
        if len(frame_stamps) < 2:
            self.velocity.update({frame_stamps[0]: PointXYZ(*INVALID_MOTION)})
            self.acceleration.update({frame_stamps[0]: PointXYZ(*INVALID_MOTION)})
            return
        # calculate reference map once for all
        reference_map = {stamp: [] for stamp in frame_stamps}
        for n, stamp in enumerate(frame_stamps[: -1]):
            for stamp_ in frame_stamps[n + 1:]:
                if abs(stamp - stamp_) < ca_interval_us_int:
                    reference_map[stamp].append(stamp_)
                    reference_map[stamp_].append(stamp)
        # svd fit of velocity
        mat_A_list = list()
        mat_b_list = list()
        for stamp, ref_stamps in reference_map.items():
            ''' Explanation of svd problem in constant-acceleration model
            t * v_x + t^2/2 * a_x = pt_x - p0_x
            t * v_y + t^2/2 * a_y = pt_y - p0_y
            Then:
            A:[t, t^2/2] * X:[[v_x, v_y], [a_x, a_y]] = b:[pt_x - p0_x, pt_y - p0_y]
            '''
            mat_A, mat_b = self.__cal_motion_prob(stamp, ref_stamps)
            mat_A_list.append(mat_A)
            mat_b_list.append(mat_b)
        motion_mat = np.linalg.pinv(diag_stack(*mat_A_list)).dot(diag_stack(*mat_b_list))
        for i, (stamp, ref_stamps) in enumerate(reference_map.items()):
            # return invalid values for no ref
            if len(ref_stamps) == 0:
                self.velocity.update({stamp: PointXYZ(*INVALID_MOTION)})
            else:
                self.velocity.update({stamp: PointXYZ(motion_mat[i*2, i*2], motion_mat[i*2, i*2+1], 0)})
        # print(f'INFO: Object {self.get_id()} velocity inferred')
        # debug_svd()
        assert set(self.traces.keys()) == set(self.velocity.keys())
        # svd fit of accleration
        mat_A_list = list()
        mat_b_list = list()
        for stamp, ref_stamps in reference_map.items():
            '''
            t * a_x = vt_x - v0_x, t * a_y = vt_y - v0_y
            Then:
            A:[t] * X:[a_x, a_y] = [vt_x - v0_x, vt_y - v0_y]
            '''
            mat_A, mat_b = self.__cal_accel_prob(stamp, ref_stamps)
            mat_A_list.append(mat_A)
            mat_b_list.append(mat_b)
        accel_mat = np.linalg.pinv(diag_stack(*mat_A_list)).dot(diag_stack(*mat_b_list))
        for i, (stamp, ref_stamps) in enumerate(reference_map.items()):
            # return invalid values for no ref
            if len(ref_stamps) == 0:
                self.acceleration.update({stamp: PointXYZ(*INVALID_MOTION)})
            else:
                self.acceleration.update({stamp: PointXYZ(accel_mat[i, i*2], accel_mat[i, i*2+1], 0)})
        # print(f'INFO: Object {self.get_id()} acceleration inferred')
        assert set(self.traces.keys()) == set(self.acceleration.keys())

    def __cal_motion_prob(self, stamp: int, ref_stamps: list) -> Tuple[NDArray, NDArray]:
        if len(ref_stamps) == 0:
            return np.ones((2, 2)), np.zeros((2, 2))
        mat_A, mat_b = list(), list()
        for ref in ref_stamps:
            mat_b.append([self.traces[ref].box.center.x - self.traces[stamp].box.center.x,
                          self.traces[ref].box.center.y - self.traces[stamp].box.center.y])
            t_gap = (ref - stamp) / STAMP_PER_SEC
            # mat_A.append([t_gap, t_gap ** 2 / 2]) disable accel in vel svd
            mat_A.append([t_gap, 0])
        return np.array(mat_A), np.array(mat_b)

    def __cal_accel_prob(self, stamp: int, ref_stamps: list) -> Tuple[NDArray, NDArray]:
        if len(ref_stamps) == 0:
            return np.ones((1, 1)), np.zeros((1, 2))
        mat_A, mat_b = list(), list()
        for ref in ref_stamps:
            mat_b.append([self.velocity[ref].x - self.velocity[stamp].x,
                          self.velocity[ref].y - self.velocity[stamp].y])
            mat_A.append([(ref - stamp) / STAMP_PER_SEC])
        return np.array(mat_A), np.array(mat_b)

    def matplot(self, dpi=200) -> None:
        assert set(self.velocity.keys()) == set(self.acceleration.keys())
        import matplotlib.pyplot as plt
        fig, subs = plt.subplots(nrows=2, ncols=2, dpi=dpi)
        fig.suptitle(self.get_id())
        subs[0, 0].set(title='x velocity')
        subs[0, 1].set(title='y velocity')
        subs[1, 0].set(title='x acceleration')
        subs[1, 1].set(title='y acceleration')
        t_seq, v_x, v_y, a_x, a_y = list(), list(), list(), list(), list()
        for stamp in self.velocity.keys():
            t_seq.append(stamp)
            v_x.append(self.velocity[stamp].x)
            v_y.append(self.velocity[stamp].y)
            a_x.append(self.acceleration[stamp].x)
            a_y.append(self.acceleration[stamp].y)
        subs[0, 0].plot(t_seq, v_x, '.-')
        subs[0, 1].plot(t_seq, v_y, '.-')
        subs[1, 0].plot(t_seq, a_x, '.-')
        subs[1, 1].plot(t_seq, a_y, '.-')
